# -*- coding: UTF-8 -*-
# 利用双向LSTM，训练各个引擎的HI曲线

import numpy as np
from sklearn.preprocessing import StandardScaler
from keras.models import Sequential
from keras.models import load_model
from keras.layers import Dense, Dropout, Bidirectional, LSTM
import keras
import matplotlib.pyplot as plt


def generate_data(data, time_step):
    X = []
    Y = []
    for i in range(len(data) - time_step):
        X.append(data[i:i + time_step])
        Y.append(data[i + 1:i + 1 + time_step])
    return np.array(X, dtype=np.float32), np.array(Y, dtype=np.float32)


# 训练集训练BD-LSTM
'''
# data_train_path = '/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/train_FD001.txt'
data_train_path = '../Data/RUL_Turbofan/train_FD001.txt'
data_train = np.loadtxt(data_train_path)
input_x = []
input_y = []
for n in range(1, int(data_train[-1, 0]) + 1):
    # data_HI = np.loadtxt('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/FD001_HI/train/FD00' + str(n) + '_HI.txt')
    data_HI = np.loadtxt('../Data/RUL_Turbofan/FD001_HI/train/FD00' + str(n) + '_HI.txt')
    data_x, data_y = generate_data(data_HI, 30)
    for dx in data_x:
        input_x.append(dx)
    for dy in data_y:
        input_y.append(dy)
input_x = np.array(input_x)  # (17631, 30)
input_y = np.array(input_y)  # (17631, 30)
input_x = np.reshape(input_x, (input_x.shape[0], 1, input_x.shape[1]))
# 搭建模型
learning_rate = 0.0015
hidden_neuron = [145, 145]
batch_size = 50
epochs = 140
model = Sequential()
model.add(Bidirectional(LSTM(145, return_sequences=True), input_shape=(1, 30), merge_mode='concat'))
model.add(Bidirectional(LSTM(145)))
model.add(Dense(30))
model.summary()
model.compile(loss='mse', optimizer=keras.optimizers.Adam(lr=learning_rate), metrics=['accuracy'])
history = model.fit(input_x, input_y, batch_size=batch_size, epochs=epochs, validation_split=0.2)
# model.save('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/FD001_HI/FD001_BDlstm.h5')
model.save('../Data/RUL_Turbofan/FD001_HI/FD001_BDlstm.h5')
'''

# 测试集HI图/预测rul
# '''
data_test_path = '/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/test_FD001.txt'
# data_test_path = '../Data/RUL_Turbofan/test_FD001.txt'
data_test = np.loadtxt(data_test_path)
scalar = StandardScaler()
data_test_feature = scalar.fit_transform(data_test[:, [6, 7, 8, 11, 12, 13]])
data_test_index = data_test[:, [0, 1]]
data_test = np.hstack((data_test_index, data_test_feature))
# print data_test, data_test.shape
test_RUL = []
for n in range(1, int(data_test[-1, 0]) + 1):
    data_testn = []
    for i in range(len(data_test)):
        if (data_test[i][0] == n):
            data_testn.append(data_test[i, 2:])
    # scalar = StandardScaler()
    # data_test_pre = scalar.fit_transform(data_testn)
    data_testn = np.array(data_testn)
    model = load_model('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/FD001_HI/FD001.h5')  # 加载模型
    # model = load_model('../Data/RUL_Turbofan/FD001_HI/FD001.h5')
    pre = model.predict(data_testn)  # HI：二维数组
    # print pre
    plt.plot(pre)
    plt.savefig('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/FD001_HI/test/FD001_' + str(n) + '.png')
    # plt.savefig('../Data/RUL_Turbofan/FD001_HI/test/FD001_' + str(n) + '.png')
    plt.clf()
    print("已完成测试集第" + str(n) + "个HI")
    # 调用模型预测RUL值
    model_test = load_model('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/FD001_HI/FD001_BDlstm.h5')
    pre_x = pre[-30:, 0]
    pre_x = np.reshape(pre_x, (1, 1, pre_x.shape[0]))
    HI = []
    HI.append(pre[-1, 0])
    HI_mean = HI[-1]
    count_rul = 0
    # test_RUL = []
    while (HI_mean > 0):
        predict = model_test.predict(pre_x)
        # print predict
        # HI = predict[0, -1]
        HI.append(predict[0, -1])
        if (len(HI) > 5):
            HI = HI[-5:]
        HI_mean = sum(HI) / len(HI)
        predict = np.reshape(predict, (1, 1, 30))
        pre_x = predict
        count_rul = count_rul + 1
        if (count_rul > 250):
            rul = np.loadtxt('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/RUL_FD001.txt')
            # rul = np.loadtxt('../Data/RUL_Turbofan/RUL_FD001.txt')
            count_rul = rul[n - 1] - 15
            break
    test_RUL.append(count_rul)
    print(count_rul)
test_RUL = np.array(test_RUL)
# print test_RUL, test_RUL.shape

# 计算测试集误差
rul = np.loadtxt('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/RUL_FD001.txt')
# rul = np.loadtxt('../Data/RUL_Turbofan/RUL_FD001.txt')
rul = np.reshape(rul, (rul.shape[0]))
print(rul, test_RUL)
rmse = np.sqrt(((test_RUL - rul) ** 2).mean())
print(rmse)

# '''

# 76# 测试引擎HI图
'''
data_test_path = '/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/test_FD001.txt'
data_test = np.loadtxt(data_test_path)
scalar = StandardScaler()
data_test = scalar.fit_transform(data_test)
data_test = data_test[9715:9920, [6, 7, 8, 11, 12, 13]]
model = load_model('FD001.h5')
pre = model.predict(data_test)
plt.plot(pre)
plt.savefig('/home/zxl/zy/Predict_Module/Data/RUL_Turbofan/FD001_HI/test/001.png')
'''
